#include "mrf.h"

/**
 * Currently for Grayscale only
 */

MRF::MRF(void)
{
	(this->means).resize((int)NUM_CLUSTERS);
	(this->variances).resize((int)NUM_CLUSTERS);
}

MRF::~MRF(void)
{
	//cvReleaseMat(&this->labelMatrix);
}

/**
 * Initiate MRF algorithm for new image with learning
 */
void MRF::initLearningMRF(IplImage *img)
{
	(this->img) = img;

	//Img info
	this->height = img->height;
	this->width = img->width;
	this->step       = img->widthStep/sizeof(uchar);
	this->channels   = img->nChannels;
	this->data    = (uchar *)img->imageData;

	this->labelMatrix = cvCreateMat(this->height, this->width, CV_32FC1);

	//this->labelImage = cvCreateImage( cvGetSize(img), 8, 1 ); 

	this->defaultLabel();

	this->calculateMRFWithLearning();
}

/**
 * Initiate MRF algorithm for new image without learning
 */
void MRF::initFastMRF(IplImage *img, IplImage *labelImg)
{
	(this->img) = img;

	//Img info
	this->height = img->height;
	this->width = img->width;
	this->step       = img->widthStep/sizeof(uchar);
	this->channels   = img->nChannels;
	this->data    = (uchar *)img->imageData;

	//this->labelMatrix = cvCreateMat(this->height, this->width, CV_32FC1);
	//this->imageLabelFromEM(labelImg);
	this->labelData = (uchar *)labelImg->imageData;
	this->labelStep = labelImg->widthStep;
	this->calculateMRF();
}

/**
 * Create default/random labels
 */
void MRF::defaultLabel()
{
	for (int i=0; i<this->height; i++)
	{
		for (int j=0; j<this->width; j++)
		{
			cvmSet(this->labelMatrix,i,j,(rand()%NUM_CLUSTERS));
		}
	}
}

/**
 * Create label matrix from EM label image
 */
//void MRF::imageLabelFromEM(IplImage *labelImg)
//{
//	int pixel;
//	int step = labelImg->widthStep;
//
//	uchar* labelData    = (uchar *)labelImg->imageData;
//
//	for (int i=0; i<labelImg->height; i++)
//	{
//		for (int j=0; j<labelImg->width; j++)
//		{
//			pixel = (int)labelData[i*step+j/**labelImg->channels*//*+k*/];
//			cvmSet(this->labelMatrix, i, j, pixel/*(pixel > 160 ? 1 : 0)*/ );// Replace pixel condition to just a pixel
//
//		}
//	}
//}

/**
 * Count means and variances for each label group
 */
void MRF::countMeansAndVariances()
{
	// Init vars
	vector<float> sum(NUM_CLUSTERS,0);
	vector<float> sumsqr(NUM_CLUSTERS,0);
	vector<float> count(NUM_CLUSTERS,0);
	int R,G,B;

	int c;
	int pixel;
	float pixel3;

	//Count Miu/Var
	for (int i=0; i<this->height; i++)
	{
		for (int j=0; j<this->width; j++)
		{
			//c = (int)cvmGet(this->labelMatrix,i,j);
			c = (int)labelData[i*this->labelStep+j];
			if(COLOR_CHANNEL_NUMBER == 3)
			{
				B = (int)this->data[i*this->step+j*this->channels+0];
				G = (int)this->data[i*this->step+j*this->channels+1];
				R = (int)this->data[i*this->step+j*this->channels+2];
				pixel3 = this->getColorSpaceAvg(R,B,G);
				sum[c] += pixel3;
				sumsqr[c] += pixel3*pixel3;
			}
			else
			{
				pixel = (int)this->data[i*this->step+j*this->channels];
				sum[c] += pixel;
				sumsqr[c] += pixel*pixel;
			}
			count[c]++;
		}
	}
	for (int l=0; l<NUM_CLUSTERS; l++)
	{
		(this->means)[l] = sum[l]/count[c];
		(this->variances)[l] = sumsqr[l]/count[c]-((this->means)[l]*(this->means)[l]);
		//(this->means)[l] = sum[l]/(width*height);
		//(this->variances)[l] = sumsqr[l]/(this->width*this->height)-((this->means)[l]*(this->means)[l]);
		////(this->variances)[l] = (sumsqr[l]-(sum[l]*sum[l])/(width*height))/(width*height);
	}
}

void MRF::setMeansAndVariances(const CvMat *means, const CvMat **variances) {
	int l=0;
	float m,v;
	double c;

	for (int i=0;i<means->rows;i++) {
		for (int j=0;j<means->cols;j++) {
			v = -1;
			m = cvmGet(means,i,j);

			for (int iv=0;iv<variances[l]->rows;iv++) {
				for (int jv=0;jv<variances[l]->cols;jv++) {
					c = cvmGet(variances[l],iv,jv);
					if(c)
					{
						v = c;
						break;
					}
				}
				if(c)
					break;
			}
			//v = cvmGet(variances[l],i,j);

			(this->means)[l] = m;
			(this->variances)[l] = v;
			l++;

		}
		cout<<endl;
	}
	
}

/**
 * Do cycles/iterations and get MRF
 * This is a learning method! Miu and Var should be given already.
 */
void MRF::calculateMRFWithLearning()
{
	//vector<int> energy(NUM_CLUSTERS,0);
	float energyAmount, energyAmountCur;
	int energyClass;

	for (int n=1; n<=ITERATION_NUMBER; n++)
	{
		this->countMeansAndVariances();

		for (int i=1; i<this->height-1; i++)
		{
			for (int j=1; j<this->width-1; j++)
			{
				energyAmountCur = energyAmount = 0;
				energyClass = -1;

				for (int l=0; l<NUM_CLUSTERS; l++)
				{
					energyAmountCur = this->calculateEnergy(l, i, j);
					if(energyClass == -1 || energyAmountCur < energyAmount)
					{
						energyAmount = energyAmountCur;
						energyClass = l;
					}

				}
				// get min energy
				cvmSet(this->labelMatrix,i,j,(double)energyClass);
				//labelData[i*this->labelStep+j] = energyClass;
			}
		}
	}
}


/**
 * Do cycles/iterations and get MRF
 * This is a learning method! Miu and Var should be given already.
 */
void MRF::calculateMRF()
{
	//vector<int> energy(NUM_CLUSTERS,0);
	float energyAmount, energyAmountCur;
	int energyClass;

	for (int n=1; n<=ITERATION_NUMBER; n++)
	{
		for (int i=1; i<this->height-1; i++)
		{
			for (int j=1; j<this->width-1; j++)
			{
				energyAmountCur = energyAmount = 0;
				energyClass = -1;

				for (int l=0; l<NUM_CLUSTERS; l++)
				{
					energyAmountCur = this->calculateEnergy(l, i, j);
					if(energyClass == -1 || energyAmountCur < energyAmount)
					{
						energyAmount = energyAmountCur;
						energyClass = l;
					}

				}
				// get min energy
				//cvmSet(this->labelMatrix,i,j,(double)energyClass);
				labelData[i*this->labelStep+j] = energyClass;
			}
		}
	}
}

float MRF::calculateEnergy(int energyClass, int i, int j)
{
	float singleton = 0;
	float pixel;
	float gibbsEnergy = calculateGibbsEnergy(energyClass, i, j);// aka Doubleton
	int R,G,B;

	if(COLOR_CHANNEL_NUMBER == 3)
	{
		B = (int)this->data[i*this->step+j*this->channels+0];
		G = (int)this->data[i*this->step+j*this->channels+1];
		R = (int)this->data[i*this->step+j*this->channels+2];
		pixel = this->getColorSpaceAvg(R, G, B);
	}
	else
	{
		pixel = (int)this->data[i*this->step+j*this->channels];
	}
	
	singleton = (pow((pixel - (this->means)[energyClass]),2)/(2*pow((this->variances)[energyClass],2))) + (1/2*log(2*3.141592653589793*pow((this->variances)[energyClass],2)));
		//singleton += pow((pixel - (this->means)[chan][energyClass]),2)/pow((this->variances)[chan][energyClass],2);// + (1/2*log(2*3.141592653589793*pow((this->variances)[energyClass],2)));
		//singleton = ((pixel - (this->means)[energyClass])/(this->variances)[energyClass]) + (1/2*log(2*3.141592653589793*pow((this->variances)[energyClass],2)));

	
	return singleton+gibbsEnergy;
}

float MRF::calculateGibbsEnergy(int energyClass, int i, int j)
{
	float energy;
	vector<int> neighbours(4);
	//neighbours[0] = (int)cvmGet(this->labelMatrix,i-1,j);
	//neighbours[1] = (int)cvmGet(this->labelMatrix,i+1,j);
	//neighbours[2] = (int)cvmGet(this->labelMatrix,i,j-1);
	//neighbours[3] = (int)cvmGet(this->labelMatrix,i,j+1);
	neighbours[0] = (int)labelData[(i-1)*this->labelStep+j];
	neighbours[1] = (int)labelData[(i+1)*this->labelStep+j];
	neighbours[2] = (int)labelData[i*this->labelStep+(j-1)];
	neighbours[3] = (int)labelData[i*this->labelStep+(j+1)];

	energy = 0;
	for(int l=0; l<4; l++)
	{
		if(energyClass == neighbours[l]) energy--;
		else energy++;
	}

    energy = MRF_BETA*energy;

	return (float)energy;
}

/**
 * Color space I3
 */
vector<int> MRF::getColorSpace(vector<int> I)
{
	//float I1, I2, I3;
	//I1 = (I[0] + I[1] + I[2])/3; // (R+G+B)/3
	//I2 = I[2]-I[0];// R-B
	//I3 = (2*I[1]-I[2]-I[0]) / 2;// (2G-R-B)/2

	//I[0] = I1;
	//I[1] = I2;
	//I[2] = I3;
	
	return I;
}

/**
 * Color space I3 average
 */
float MRF::getColorSpaceAvg(int R, int G, int B)
{
	float I1, I2, I3;
	I1 = (R + G + B)/3; // (R+G+B)/3
	I2 = R-B;// R-B
	I3 = (2*G-R-B) / 2; // (2G-R-B)/2
	
	//return ((I1+I2+I3)/3);
	return ((R+G+B)/3);
	//return (R+G)/2;//HSV
}

/**
 * Outputs. Mostly for debuging
 */
void MRF::ouputLabelledImage()
{
	float c;
	char key;
	IplImage* outImg = cvCreateImage( cvGetSize(this->img), 8, 1 );

	if(COLOR_CHANNEL_NUMBER > 1) cvCvtColor(this->img,outImg,CV_BGR2GRAY);
	else outImg = this->img;

	for (int i=1; i<this->height-1; i++)
	{
		for (int j=1; j<this->width-1; j++)
		{
			c = cvmGet(this->labelMatrix,i,j);
			(outImg->imageData + i*outImg->widthStep)[j] = c*50;
			//cout<<c<<endl;
		}
	}
	cvNamedWindow( "out", 1 );
	for(;;)
    {
		cvShowImage( "out", outImg );
	    
		key =(char) cvWaitKey(0);
		if( key == 27 || key == 'q' || key == 'Q' ) // 'ESC'
			break;
	}
	cvDestroyWindow( "out" );

}

/**
 * Returns labelled image to main program.
 */
IplImage* MRF::getLabelledImage()
{
	float c;
	char key;
	int offset = 255 / NUM_CLUSTERS;
	IplImage* outImg = cvCreateImage( cvGetSize(this->img), 8, 1 );

	if(COLOR_CHANNEL_NUMBER > 1) cvCvtColor(this->img,outImg,CV_BGR2GRAY);
	else outImg = this->img;

	for (int i=1; i<this->height-1; i++)
	{
		for (int j=1; j<this->width-1; j++)
		{
			//c = cvmGet(this->labelMatrix,i,j);
			c = labelData[i*this->labelStep+j];
			(outImg->imageData + i*outImg->widthStep)[j] = c*offset;
		}
	}

	return outImg;
	//return this->labelImage;

}